import torch
import torch.nn as nn
import torch.utils.data
import torch.utils.data.distributed
from elderlyCare.HumanFAT.SPPE.src.utils.img import flip, shuffleLR
from elderlyCare.HumanFAT.SPPE.src.models.FastPose import FastPose

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
import sys
import torch._utils


try:
    torch._utils._rebuild_tensor_v2
except AttributeError:
    def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks):
        tensor = torch._utils._rebuild_tensor(storage, storage_offset, size, stride)
        tensor.requires_grad = requires_grad
        tensor._backward_hooks = backward_hooks
        return tensor
    torch._utils._rebuild_tensor_v2 = _rebuild_tensor_v2


class InferenNet(nn.Module):
    """
    InferenNet
    """

    # constants
    WEIGHT_FILE_RES101 = "D:\\Develop\\pycharmWorkspace\\Intelligent-Endowment\\generate\\models\\sppe\\fast_res101_320x256.pth"

    def __init__(self, dataset, weights_file=WEIGHT_FILE_RES101):
        super().__init__()

        self.pyranet = FastPose('resnet101').to(device)
        print('Loading pose model from {}'.format(weights_file))
        sys.stdout.flush()
        self.pyranet.load_state_dict(torch.load(weights_file, map_location='cpu'))
        self.pyranet.eval()
        self.pyranet = model

        self.dataset = dataset

    def forward(self, x):
        out = self.pyranet(x)
        out = out.narrow(1, 0, 17)

        flip_out = self.pyranet(flip(x))
        flip_out = flip_out.narrow(1, 0, 17)

        flip_out = flip(shuffleLR(
            flip_out, self.dataset))

        out = (flip_out + out) / 2

        return out


class InferenNet_fast(nn.Module):
    """
    InferenNet_fast with res101.
    """

    # constants
    WEIGHT_FILE_RES101 = "D:\\Develop\\pycharmWorkspace\\Intelligent-Endowment\\generate\\models\\sppe\\fast_res101_320x256.pth"

    def __init__(self, weights_file=WEIGHT_FILE_RES101):
        super().__init__()

        self.pyranet = FastPose('resnet101').cuda()
        print('Loading pose model from {}'.format(weights_file))
        self.pyranet.load_state_dict(torch.load(weights_file, map_location='cpu'))
        self.pyranet.eval()

    def forward(self, x):
        out = self.pyranet(x)
        out = out.narrow(1, 0, 17)

        return out


class InferenNet_fastRes50(nn.Module):
    """
    InferenNet_fast with res50.
    """

    # constants
    WEIGHT_FILE_RES50 = "D:\\Develop\\pycharmWorkspace\\Intelligent-Endowment\\generate\\models\\sppe\\fast_res50_256x192.pth"

    def __init__(self, weights_file=WEIGHT_FILE_RES50):
        super().__init__()

        self.pyranet = FastPose('resnet50', 17).to(device)
        print('Loading pose model from {}'.format(weights_file))
        self.pyranet.load_state_dict(torch.load(weights_file, map_location='cpu'))
        self.pyranet.eval()

    def forward(self, x):
        out = self.pyranet(x)

        return out
